# 2009.09.16 chengtao 
# import OPLS force field parameter from GROMACS OPLS data base

#constants 
BHARM = 836.8
AHARM = 8.368
TCOSP = 8.368
N12_6_SIGMA = 0.11246
N12_6_EPSILON = 4.184

def RB2Fourier(plist):
    """ Ryckaert-Bellemans function to Ryckaert-Bellemans function
    reference GROMACS mannual"""
    rb = [float(i) for i in plist]
    f = [0,0,0]
    f[0] = -2*rb[1] -1.5*rb[3]
    f[1] = -rb[4] - rb[2]
    f[2] = -0.5 * rb[3]
    f = [ i/TCOSP for i in f]
    counter = 0.0
    parameter =""
    parameter += "%5.1f,%10.4f,%5.1f"%(0.0, f[0], 1.0)
    parameter += ",%5.1f,%10.4f,%5.1f"%(180.0, f[1], 2.0)
    parameter += ",%5.1f,%10.4f,%5.1f"%(0.0, f[2], 3.0)
    return parameter
def readGromacsFF():
    """import force field parameters from gromacs. 
    nonBond: nonbond parameter including atomtype eqt type charge and vdw
    atomType: detail explaination of the atomtype
    bond: bond parameters"""
    f = open("ffoplsaanb.itp", 'r')
    nonBond = {}
    for i in f:
        if i.startswith(";") or i.startswith("#") or i.startswith("["):
            pass
        else:
            # name  bond_type    mass    charge   ptype          sigma      epsilon
            nonBond[i.split()[0]] = [j for j in i.split(';')[0].split()]
    f.close()
    f = open("ffoplsaa.atp", 'r')
    atomType = {}
    for i in f:
        if i.startswith(";"):
            pass
        else:
            atomType[i.split()[0]] = i.split(';')[1]
    f.close()
    f = open("ffoplsaabon.itp", 'r')
    bond = []
    for i in f:
        if i.startswith("[") or i.startswith(";") or i.strip()== "" or i.startswith("#"):
            pass
        else:
            bond.append(i)
    f.close()
    return nonBond, atomType, bond
    
def writeEQT(nonBond): 
    """write eqt file according the information in nonBond (fflopsaanb.itp)"""
    o = open("opls.eqt", 'w')
    o.write("# AAT      :   NB        ATC          BINC  Bond  A/C   A/S   T/C   T/S   O/C   O/S\n")
    keys = [i for i in nonBond.keys()]
    keys.sort()
    for i in keys:
        o.write("%10s : %10s%10s"%(i, i, i))
        j = nonBond[i]
        for iter in range(8):
            o.write("%6s"%j[1])
        o.write("\n")
    o.close()

def wirtePPF(bond):
    o = open("opls.ppf", 'w')
    o.write("#DFF:PPF\n#PROTOCOL = AMBER\n")
    outfile =[]
    #bond parameter
    for i in bond:
        #bond
        if len(i.split(";")[0].split()) == 5:
            #  OW    HW      1    0.09572   502080.0   ; For TIP4F Water - wlj 1/98
            if len(i.split(";")) > 1:        
                ref = i.split(";",1)[1]
            else:
                ref = "\n"
            p = i.split(";")[0].split()
            bond = float(p[3]) * 10
            force = float(p[4]) / BHARM
            #BHARM: c_4, c_4:     1.5246*,   220.0742*
            line = "BHARM:%10s,%10s:%10.4f,%10.4f:%s"%(p[0], p[1], bond, force, ref)
            outfile.append(line)
        #angle
        elif len(i.split(";")[0].split()) == 6:
            if len(i.split(";")) > 1:        
                ref = i.split(";",1)[1]
            else:
                ref = "\n"
            p = i.split(";")[0].split()
            bond = float(p[4]) * 10
            force = float(p[5]) / AHARM
            line = "AHARM:%10s,%10s,%10s:%10.4f,%10.4f:%s"%(p[0], p[1], p[2], bond, force, ref)
            outfile.append(line)
        #torsion
        elif len(i.split(";")[0].split()) == 11:
            if len(i.split(";")) > 1:        
                ref = i.split(";",1)[1]
            else:
                ref = "\n"
            p = i.split(";")[0].split()
            tcosp = RB2Fourier(p[5:11])
            if tcosp == "":
                pass
            else:
                line = "TCOSP:%10s,%10s,%10s,%10s:%68s:%s"%(p[0], p[1], p[2],p[3], tcosp, ref)
                outfile.append(line)
    
    #ATC
    for i in nonBond.keys():
        line = "ATC:%10s:%10.4f\n"%(i, float(nonBond[i][4]))
        outfile.append(line)
    #N12_6
    for i in nonBond.keys():
        sigma = float(nonBond[i][6]) / N12_6_SIGMA
        epsilon = float(nonBond[i][6]) / N12_6_EPSILON
        line ="N12_6: %10s:%10.4f,%10.4f\n"%(i, sigma, epsilon)
        outfile.append(line)
    
    outfile.sort()        
    for i in outfile:
        o.write(i)
    o.close()
if __name__ == "__main__":
    nonBond, atomType, bond = readGromacsFF()
    writeEQT(nonBond)
    wirtePPF(bond)